import os
import ujson
import aicube
import time
import math
from machine import FPIOA, I2C
from libs.PipeLine import ScopedTiming
from libs.Utils import *
from media.sensor import *
from media.display import *
from media.media import *
import nncase_runtime as nn
import ulab.numpy as np
import image
import gc

# ================== 系统配置 ==================
display_mode = "lcd"
if display_mode == "lcd":
    DISPLAY_WIDTH = ALIGN_UP(800, 16)
    DISPLAY_HEIGHT = 480
else:
    DISPLAY_WIDTH = ALIGN_UP(1920, 16)
    DISPLAY_HEIGHT = 1080

OUT_RGB888P_WIDTH = ALIGN_UP(1280, 16)
OUT_RGB888P_HEIGH = 720

root_path = "/sdcard/mp_deployment_source/"
config_path = root_path + "deploy_config.json"
debug_mode = 1

# ================== BMP280配置 ==================
BMP280_ADDR = 0x76
bmp280_present = False
bmp280_frame_count = 0
bmp280_read_interval = 60  # 降低读取频率，每60帧读取一次
current_frame = 0
bmp280_i2c = None  # 全局I2C对象
last_valid_pressure = None
last_valid_temperature = None
error_count = 0
MAX_ERROR_COUNT = 10  # 最大连续错误次数

# BMP280校准参数
digT1 = 0
digT2 = 0
digT3 = 0
digP1 = 0
digP2 = 0
digP3 = 0
digP4 = 0
digP5 = 0
digP6 = 0
digP7 = 0
digP8 = 0
digP9 = 0

# ================== 污渍类型配置 ==================
STAIN_PRIORITY = {
    "Droppings": 1,
    "Leaves": 2,
    "Branch": 3,
    "Paper": 4
}

CLEANING_ADVICE = {
    "Droppings": "立即清洁！具有腐蚀性",
    "Leaves": "尽快清理，避免堆积",
    "Branch": "中等优先级，可安排清理",
    "Paper": "低优先级，最后处理"
}

PRIORITY_COLORS = {
    1: (255, 0, 0),
    2: (255, 165, 0),
    3: (255, 255, 0),
    4: (0, 255, 0)
}

def two_side_pad_param(input_size, output_size):
    ratio_w = output_size[0] / input_size[0]
    ratio_h = output_size[1] / input_size[1]
    ratio = min(ratio_w, ratio_h)
    new_w = int(ratio * input_size[0])
    new_h = int(ratio * input_size[1])
    dw = (output_size[0] - new_w) / 2
    dh = (output_size[1] - new_h) / 2
    top = int(round(dh - 0.1))
    bottom = int(round(dh + 0.1))
    left = int(round(dw - 0.1))
    right = int(round(dw - 0.1))
    return top, bottom, left, right, ratio

def read_deploy_config(config_path):
    with open(config_path, 'r') as f:
        try:
            return ujson.load(f)
        except Exception as e:
            print(f"配置文件解析错误: {e}")
            return {}

def get_priority_info(det_boxes, labels):
    priority_counts = {1: 0, 2: 0, 3: 0, 4: 0}
    highest_priority = 5
    highest_priority_stain = "无污渍"

    for box in det_boxes:
        stain_type = labels[box[0]]
        priority = STAIN_PRIORITY.get(stain_type, 4)

        if priority < highest_priority:
            highest_priority = priority
            highest_priority_stain = stain_type

        if priority in priority_counts:
            priority_counts[priority] += 1

    return highest_priority, highest_priority_stain, priority_counts

# ================== BMP280函数 ==================
def bmp280_init():
    global digT1, digT2, digT3, digP1, digP2, digP3, digP4, digP5, digP6, digP7, digP8, digP9
    global bmp280_present, bmp280_i2c, error_count

    # 重置错误计数器
    error_count = 0

    # 配置I2C引脚
    fpioa = FPIOA()
    fpioa.set_function(44, FPIOA.IIC3_SCL)
    fpioa.set_function(45, FPIOA.IIC3_SDA)

    try:
        # 创建全局I2C对象（如果不存在）
        if bmp280_i2c is None:
            bmp280_i2c = I2C(3, freq=400_000)
            print("创建新的I2C对象")

        # 扫描I2C设备
        devices = bmp280_i2c.scan()
        print("检测到I2C设备:", [hex(addr) for addr in devices])

        if BMP280_ADDR not in devices:
            print(f"错误: 在I2C地址0x{BMP280_ADDR:02x}未找到BMP280")
            print("请检查: 1) 传感器接线 2) 供电 3) I2C地址配置")
            bmp280_present = False
            return False

        # 读取校准参数
        calib = bmp280_i2c.readfrom_mem(BMP280_ADDR, 0x88, 24)

        # 解析温度参数
        digT1 = (calib[1] << 8) | calib[0]
        digT2 = (calib[3] << 8) | calib[2]
        if digT2 > 32767:
            digT2 -= 65536
        digT3 = (calib[5] << 8) | calib[4]
        if digT3 > 32767:
            digT3 -= 65536

        # 解析气压参数
        digP1 = (calib[7] << 8) | calib[6]
        digP2 = (calib[9] << 8) | calib[8]
        if digP2 > 32767:
            digP2 -= 65536
        digP3 = (calib[11] << 8) | calib[10]
        if digP3 > 32767:
            digP3 -= 65536
        digP4 = (calib[13] << 8) | calib[12]
        if digP4 > 32767:
            digP4 -= 65536
        digP5 = (calib[15] << 8) | calib[14]
        if digP5 > 32767:
            digP5 -= 65536
        digP6 = (calib[17] << 8) | calib[16]
        if digP6 > 32767:
            digP6 -= 65536
        digP7 = (calib[19] << 8) | calib[18]
        if digP7 > 32767:
            digP7 -= 65536
        digP8 = (calib[21] << 8) | calib[20]
        if digP8 > 32767:
            digP8 -= 65536
        digP9 = (calib[23] << 8) | calib[22]
        if digP9 > 32767:
            digP9 -= 65536

        # 设置控制寄存器
        bmp280_i2c.writeto_mem(BMP280_ADDR, 0xF4, b'\x27')
        time.sleep_ms(10)

        print("BMP280初始化成功")
        bmp280_present = True
        return True
    except Exception as e:
        print(f"BMP280初始化失败: {e}")
        # 尝试重新创建I2C对象
        try:
            bmp280_i2c = I2C(3, freq=400_000)
            print("重建I2C对象")
        except:
            print("无法重建I2C对象")

        bmp280_present = False
        return False

def read_bmp280():
    global bmp280_frame_count, last_valid_pressure, last_valid_temperature, error_count, bmp280_present

    if not bmp280_present or bmp280_i2c is None:
        return last_valid_pressure, last_valid_temperature

    try:
        # 读取原始数据
        data = bmp280_i2c.readfrom_mem(BMP280_ADDR, 0xF7, 6)

        # 解析原始值
        press_raw = (data[0] << 12) | (data[1] << 4) | (data[2] >> 4)
        temp_raw = (data[3] << 12) | (data[4] << 4) | (data[5] >> 4)

        # 温度补偿计算
        var1 = ((temp_raw) / 16384.0 - (digT1) / 1024.0) * (digT2)
        var2 = (((temp_raw) / 131072.0 - (digT1) / 8192.0) *
                ((temp_raw) / 131072.0 - (digT1) / 8192.0)) * (digT3)
        t_fine = var1 + var2
        temperature = t_fine / 5120.0

        # 气压补偿计算
        var1 = t_fine / 2.0 - 64000.0
        var2 = var1 * var1 * digP6 / 32768.0
        var2 = var2 + var1 * digP5 * 2.0
        var2 = var2 / 4.0 + digP4 * 65536.0
        var1 = (digP3 * var1 * var1 / 524288.0 + digP2 * var1) / 524288.0
        var1 = (1.0 + var1 / 32768.0) * digP1

        if var1 == 0:
            pressure = 0
        else:
            p = 1048576.0 - press_raw
            p = (p - (var2 / 4096.0)) * 6250.0 / var1
            var1 = digP9 * p * p / 2147483648.0
            var2 = p * digP8 / 32768.0
            pressure = p + (var1 + var2 + digP7) / 16.0

        # 转换为hPa
        pressure = pressure / 100.0

        # 检查数据有效性
        if pressure < 300 or pressure > 1100 or temperature < -40 or temperature > 85:
            raise ValueError("传感器读数超出有效范围")

        bmp280_frame_count += 1
        error_count = 0  # 重置错误计数器

        # 更新最后一次有效值
        last_valid_pressure = pressure
        last_valid_temperature = temperature

        return pressure, temperature
    except OSError as e:
        error_count += 1
        print(f"BMP280通信错误 ({error_count}/{MAX_ERROR_COUNT}): {e}")

        # 连续错误超过阈值时尝试重新初始化
        if error_count >= MAX_ERROR_COUNT:
            print("连续错误过多，尝试重新初始化传感器...")
            if bmp280_init():
                error_count = 0
            else:
                print("重新初始化失败，暂停传感器读取")
                bmp280_present = False

        return last_valid_pressure, last_valid_temperature
    except Exception as e:
        error_count += 1
        print(f"BMP280数据处理异常 ({error_count}/{MAX_ERROR_COUNT}): {e}")
        return last_valid_pressure, last_valid_temperature

def detection():
    global current_frame, bmp280_present

    print("目标检测系统启动...")

    # ================== 初始化阶段 ==================
    # 1. 加载部署配置
    deploy_conf = read_deploy_config(config_path)
    if not deploy_conf:
        print("错误：无法加载部署配置，系统退出")
        return -1

    # 从配置中获取参数
    kmodel_name = deploy_conf["kmodel_path"]
    labels = deploy_conf["categories"]
    confidence_threshold = deploy_conf["confidence_threshold"]
    nms_threshold = deploy_conf["nms_threshold"]
    img_size = deploy_conf["img_size"]
    num_classes = deploy_conf["num_classes"]
    color_four = get_colors(num_classes)
    nms_option = deploy_conf["nms_option"]
    model_type = deploy_conf["model_type"]

    if model_type == "AnchorBaseDet":
        anchors = deploy_conf["anchors"][0] + deploy_conf["anchors"][1] + deploy_conf["anchors"][2]
    else:
        anchors = []

    kmodel_frame_size = img_size
    frame_size = [OUT_RGB888P_WIDTH, OUT_RGB888P_HEIGH]
    strides = [8, 16, 32]

    # 2. 计算预处理填充参数
    top, bottom, left, right, ratio = two_side_pad_param(frame_size, kmodel_frame_size)

    # 3. 初始化AI推理引擎
    kpu = nn.kpu()
    kpu.load_kmodel(root_path + kmodel_name)

    # 4. 初始化AI2D预处理引擎
    ai2d = nn.ai2d()
    ai2d.set_dtype(nn.ai2d_format.NCHW_FMT, nn.ai2d_format.NCHW_FMT, np.uint8, np.uint8)
    ai2d.set_pad_param(True, [0, 0, 0, 0, top, bottom, left, right], 0, [114, 114, 114])
    ai2d.set_resize_param(True, nn.interp_method.tf_bilinear, nn.interp_mode.half_pixel)
    ai2d_builder = ai2d.build(
        [1, 3, OUT_RGB888P_HEIGH, OUT_RGB888P_WIDTH],
        [1, 3, kmodel_frame_size[1], kmodel_frame_size[0]]
    )

    # 5. 初始化BMP280传感器
    bmp280_ok = bmp280_init()
    print(f"BMP280初始化状态: {'成功' if bmp280_ok else '失败'}")

    # 6. 初始化摄像头传感器
    sensor = Sensor()
    sensor.reset()
    sensor.set_hmirror(False)
    sensor.set_vflip(False)

    sensor.set_framesize(width=DISPLAY_WIDTH, height=DISPLAY_HEIGHT)
    sensor.set_pixformat(PIXEL_FORMAT_YUV_SEMIPLANAR_420)

    sensor.set_framesize(width=OUT_RGB888P_WIDTH, height=OUT_RGB888P_HEIGH, chn=CAM_CHN_ID_2)
    sensor.set_pixformat(PIXEL_FORMAT_RGB_888_PLANAR, chn=CAM_CHN_ID_2)

    # 7. 配置显示系统
    sensor_bind_info = sensor.bind_info(x=0, y=0, chn=CAM_CHN_ID_0)
    Display.bind_layer(**sensor_bind_info, layer=Display.LAYER_VIDEO1)

    if display_mode == "lcd":
        Display.init(Display.ST7701, to_ide=True)
    else:
        Display.init(Display.LT9611, to_ide=True)

    osd_img = image.Image(DISPLAY_WIDTH, DISPLAY_HEIGHT, image.ARGB8888)

    # 8. 初始化媒体管理器
    MediaManager.init()

    # 9. 启动摄像头
    sensor.run()

    # 10. 初始化AI2D输出张量
    data = np.ones((1, 3, kmodel_frame_size[1], kmodel_frame_size[0]), dtype=np.uint8)
    ai2d_output_tensor = nn.from_numpy(data)

    # 初始化温度和压力值
    pressure = None
    temperature = None

    # ================== 主循环 ==================
    print("进入主循环...")
    while True:
        current_frame += 1

        with ScopedTiming("total", debug_mode > 0):
            # 按指定频率读取BMP280数据
            if bmp280_present and current_frame % bmp280_read_interval == 0:
                pressure, temperature = read_bmp280()
                if pressure is not None and temperature is not None:
                    print(f"传感器读数 - 气压: {pressure:.1f} hPa, 温度: {temperature:.1f}°C")
                else:
                    print("警告: BMP280数据无效，使用最后有效值")

            # 捕获图像
            rgb888p_img = sensor.snapshot(chn=CAM_CHN_ID_2)

            if rgb888p_img.format() == image.RGBP888:
                ai2d_input = rgb888p_img.to_numpy_ref()
                ai2d_input_tensor = nn.from_numpy(ai2d_input)

                # 执行AI2D预处理
                ai2d_builder.run(ai2d_input_tensor, ai2d_output_tensor)

                # 模型推理
                kpu.set_input_tensor(0, ai2d_output_tensor)
                kpu.run()

                # 获取输出结果
                results = []
                for i in range(kpu.outputs_size()):
                    out_data = kpu.get_output_tensor(i)
                    result = out_data.to_numpy().flatten()
                    del out_data
                    results.append(result)

                # 后处理
                det_boxes = aicube.anchorbasedet_post_process(
                    results[0], results[1], results[2],
                    kmodel_frame_size, frame_size, strides,
                    num_classes, confidence_threshold,
                    nms_threshold, anchors, nms_option
                )

                # 污渍优先级评估
                highest_priority = 5
                highest_stain = "无污渍"
                priority_counts = {1: 0, 2: 0, 3: 0, 4: 0}

                if det_boxes:
                    highest_priority, highest_stain, priority_counts = get_priority_info(det_boxes, labels)

                # 绘制UI
                osd_img.clear()

                # ========== 顶部状态栏绘制 ==========
                status_bar_height = 60
                osd_img.draw_rectangle(0, 0, DISPLAY_WIDTH, status_bar_height,
                                      color=(0, 0, 0, 180), fill=True)

                priority_color = PRIORITY_COLORS.get(highest_priority, (255, 255, 255))
                priority_text = f"清洁优先级: {highest_stain} (等级{highest_priority})"
                osd_img.draw_string_advanced(10, 10, 32, priority_text, color=priority_color)

                if highest_stain != "无污渍":
                    advice = CLEANING_ADVICE.get(highest_stain, "建议清理")
                    advice_text = f"建议: {advice}"
                    osd_img.draw_string_advanced(10, 40, 24, advice_text, color=(255, 255, 255))

                # ========== 污渍统计面板 ==========
                stats_x = DISPLAY_WIDTH - 150
                stats_y = 10
                osd_img.draw_string_advanced(stats_x, stats_y, 24, "污渍统计:", color=(200, 200, 200))
                stats_y += 30

                for prio in sorted(priority_counts.keys()):
                    count = priority_counts[prio]
                    if count > 0:
                        color = PRIORITY_COLORS.get(prio, (255, 255, 255))
                        stain_name = [k for k, v in STAIN_PRIORITY.items() if v == prio][0]
                        stat_text = f"P{prio}: {stain_name}x{count}"
                        osd_img.draw_string_advanced(stats_x, stats_y, 24, stat_text, color=color)
                        stats_y += 30

                # ========== 传感器状态面板 ==========
                sensor_x = DISPLAY_WIDTH - 200
                sensor_y = DISPLAY_HEIGHT - 90

                # 绘制背景
                osd_img.draw_rectangle(sensor_x, sensor_y, 190, 80,
                                      color=(0, 0, 0, 150), fill=True)

                # 传感器状态
                #status_text = f"传感器: {'正常' if bmp280_present else '故障'}"
                #status_color = (0, 255, 0) if bmp280_present else (255, 0, 0)
                #osd_img.draw_string_advanced(sensor_x + 10, sensor_y + 5, 20, status_text, color=status_color)

                # 显示温度
                if temperature is not None:
                    temp_text = f"温度: {temperature:.1f}°C"
                    osd_img.draw_string_advanced(sensor_x + 10, sensor_y + 25, 20, temp_text, color=(0, 200, 255))
                else:
                    osd_img.draw_string_advanced(sensor_x + 10, sensor_y + 25, 20, "温度: --", color=(200, 200, 200))

                # 显示气压
                if pressure is not None:
                    press_text = f"气压: {pressure:.1f} hPa"
                    osd_img.draw_string_advanced(sensor_x + 10, sensor_y + 45, 20, press_text, color=(0, 200, 255))
                else:
                    osd_img.draw_string_advanced(sensor_x + 10, sensor_y + 45, 20, "气压: --", color=(200, 200, 200))

                # 显示帧计数
                #frame_text = f"帧号: {current_frame}"
                #osd_img.draw_string_advanced(sensor_x + 10, sensor_y + 65, 18, frame_text, color=(180, 180, 180))

                # ========== 绘制检测框和标签 ==========
                if det_boxes:
                    for box in det_boxes:
                        x1, y1, x2, y2 = box[2], box[3], box[4], box[5]
                        stain_type = labels[box[0]]
                        confidence = box[1]
                        priority = STAIN_PRIORITY.get(stain_type, 4)
                        color = PRIORITY_COLORS.get(priority, (255, 255, 255))

                        x = int(x1 * DISPLAY_WIDTH / OUT_RGB888P_WIDTH)
                        y = int(y1 * DISPLAY_HEIGHT / OUT_RGB888P_HEIGH) + status_bar_height
                        w = int((x2 - x1) * DISPLAY_WIDTH / OUT_RGB888P_WIDTH)
                        h = int((y2 - y1) * DISPLAY_HEIGHT / OUT_RGB888P_HEIGH)

                        osd_img.draw_rectangle(x, y, w, h, color=color, thickness=3)
                        label_text = f"{stain_type} {confidence:.2f}"
                        osd_img.draw_string_advanced(x, y - 30, 28, label_text, color=color)

                # 显示结果到OSD层
                Display.show_image(osd_img, 0, 0, Display.LAYER_OSD3)

                # ========== 控制台输出 ==========
                print("\n" + "=" * 50)
                print(f"检测到污渍数量: {len(det_boxes)}")

                if det_boxes:
                    print(f"最高优先级污渍: {highest_stain} (等级{highest_priority})")
                    print("污渍分布统计:")
                    for prio in sorted(priority_counts.keys()):
                        count = priority_counts[prio]
                        if count > 0:
                            stain_name = [k for k, v in STAIN_PRIORITY.items() if v == prio][0]
                            print(f"  等级{prio}: {stain_name} - {count}个")
                    print(f"清洁建议: {CLEANING_ADVICE.get(highest_stain, '常规清理')}")
                else:
                    print("未检测到污渍，区域清洁")

                # 输出BMP280数据
                if bmp280_present:
                    if pressure is not None and temperature is not None:
                        print(f"环境数据 - 气压: {pressure:.1f} hPa, 温度: {temperature:.1f}°C")
                    else:
                        print("警告: BMP280数据无效，使用最后有效值")
                else:
                    print("传感器状态: 已禁用")

                # 定期内存回收
                if current_frame % 100 == 0:
                    gc.collect()
                    print("执行内存回收")

            rgb888p_img = None

    print("系统正在关闭...")
    sensor.stop()
    Display.deinit()
    MediaManager.deinit()
    del ai2d_output_tensor
    gc.collect()
    time.sleep(1)
    nn.shrink_memory_pool()
    print("系统已关闭")
    return 0

if __name__ == "__main__":
    detection()
